/*
 * AqueousPhase.cpp
 *
 *  Created on: 23 May 2011
 *      Author: allan
 */

#include "AqueousPhase.h"

// C++ includes
#include <algorithm>
#include <functional>
#include <numeric>

// GeoReact includes
#include "Database.h"
#include "NumericalUtilities.h"

AqueousPhase::AqueousPhase() 
{}

AqueousPhase::AqueousPhase(const vector<string>& species) :
species(species)
{
	// Setting the ionic and neutral species
	for(unsigned ispecies = 0; ispecies < species.size(); ++ispecies)
	{
		// The electrical charge of the current species
		const double z = ElectricalCharge(species[ispecies]);
		
		// Setting the indexes and charges of the ions
		if(z != 0)
		{
			ions.push_back(ispecies);
			zi.push_back(z);
		}
		// Setting the indexes of the neutral species
		else
		{
			neutrals.push_back(ispecies);
			
			if(species[ispecies] == "H2O")
				iwater = ispecies;
		}
	}
	
	InitElectrolytes();
	InitDefaultActivityModels();
}

void AqueousPhase::SetIonicActivityModel(IonicActivityModel model)
{
	ionicActivityFunc = HkfIonicActivityModel(GetIonicSpecies());
}

void AqueousPhase::SetSetschenowConstant(const string& speciesName, const vector<double>& bi, const vector<double>& Ti)
{
	// Loop over all the neutral species
	for(unsigned n = 0; n < neutrals.size(); ++n)
		if(speciesName == species[neutrals[n]])
			neutralActivityFuncs[n] = SetschenowActivityModel(bi, Ti);
}

const vector<string>& AqueousPhase::GetSpecies() const
{
	return species;
}

const vector<string> AqueousPhase::GetIonicSpecies() const
{
	vector<string> ionicSpecies(ions.size());
	
	for(unsigned i = 0; i < ions.size(); ++i)
		ionicSpecies[i] = species[ions[i]];
	
	return ionicSpecies;
}

const vector<double> AqueousPhase::GetIonicCharges() const
{
	return zi;
}

const Index AqueousPhase::GetWaterIndex() const
{
	return iwater;
}

const double AqueousPhase::EffectiveIonicStrength(const vector<double>& n) const
{
	// The effective ionic strength to be computed
	double Ie = 0.0;
	
	// Loop over all the ions
	for(unsigned i = 0; i < ions.size(); ++i)
		Ie += zi[i] * zi[i] * n[ions[i]];
	
	// The molality factor to trasnform Ie in a molality scale
	const double molalityFactor = 55.508 / n[iwater];
	
	return 0.5 * molalityFactor * Ie;
}

const double AqueousPhase::StoichiometricIonicStrength(const vector<double>& n) const
{
	// The stoichiometric ionic strength to be computed
	double Is = 0.0;
	
	// The stoichiometric molality of the ions
	const vector<double> smions = IonicStoichiometricMolalities(n); 
	
	// Loop over all the ions
	for(unsigned i = 0; i < ions.size(); ++i)
		Is += zi[i] * zi[i] * smions[i];
	
	// The molality factor to trasnform number of moles into molality
	const double molalityFactor = 55.508 / n[iwater]; 
	
	return 0.5 * molalityFactor * Is;
}

const double AqueousPhase::TotalMolalitySolutes(const vector<double>& n) const
{
	// The molality factor to trasnform number of moles into molality
	const double molalityFactor = 55.508 / n[iwater]; 
	
	// The total number of moles
	const double nT = std::accumulate(n.begin(), n.end(), 0.0);
	
	// Remember to subtract the molality of the water species, i.e. 55.508 molal
	return molalityFactor * nT - 55.508;
}

const vector<double> AqueousPhase::Molalities(const vector<double>& n) const
{
	// The molalities of the species to be computed
	vector<double> m(n);
	
	// The molality factor to trasnform number of moles into molality
	const double molalityFactor = 55.508 / n[iwater]; 
	
	// Loop over all the species
	for(unsigned i = 0; i < n.size(); ++i)
		m[i] *= molalityFactor;
		
	return m;
}

const vector<double> AqueousPhase::IonicMolalities(const vector<double>& n) const
{
	// The ionic molalities to be computed
	vector<double> mi(ions.size());
	
	// The molality factor to trasnform number of moles into molality
	const double molalityFactor = 55.508 / n[iwater]; 
	
	// Loop over all the ions
	for(unsigned i = 0; i < ions.size(); ++i)
		mi[i] = n[ions[i]] * molalityFactor;
	
	return mi;
}

const vector<double> AqueousPhase::IonicStoichiometricMolalities(const vector<double>& n) const
{
	const unsigned numIons         = ions.size();         // The number of ionic species in the aqueous phase
	const unsigned numElectrolytes = electrolytes.size(); // The number of electrolytes in the aqueous phase
	
	VectorXd mi(numIons);         // The molalities of the ions
	VectorXd me(numElectrolytes); // The molalities of the electrolytes
	
	const double molalityFactor = 55.508/n[iwater];
	
	for(unsigned i = 0; i < numIons; ++i)
		mi[i] = molalityFactor * n[ions[i]];
	
	for(unsigned e = 0; e < numElectrolytes; ++e)
		me[e] = molalityFactor * n[electrolytes[e]];
	
	VectorXd res = mi + vei.transpose() * me;
	
	return vector<double>(res.data(), res.data() + res.rows());
}

const vector<double> AqueousPhase::ActivitiyCoefficients(double T, double P, const vector<double>& n) const
{
	vector<double> gi = ActivityCoefficientsIonicSpecies(T, P, n);   // The activity coefficients of the ionic species
	vector<double> gn = ActivityCoefficientsNeutralSpecies(T, P, n); // The activity coefficients of the neutral species
	
	vector<double> g(n.size()); // The activity coefficients of the aqueous species
	
	// Loop over all the ions and set the activity coefficient vector g with the ionic activity coefficients
	for(unsigned ion = 0; ion < ions.size(); ++ion)
		g[ions[ion]] = gi[ion];
	
	// Loop over all the neutral species and set the activity coefficient vector g with the neutral activity coefficients
	for(unsigned neutral = 0; neutral < neutrals.size(); ++neutral)
		g[neutrals[neutral]] = gn[neutral];
	
	return g;
}

const vector<double> AqueousPhase::Activities(double T, double P, const vector<double>& n) const
{
	vector<double> a(ActivitiyCoefficients(T, P, n));
	
	const double molalityFactor = 55.508/n[iwater];
	
	for(unsigned i = 0; i < n.size(); ++i)
		a[i] *= n[i] * molalityFactor;
	
	return a;
}

const bool IsElectrolyte(const SpeciesInfo& speciesInfo, const ReactionInfo& reactionInfo)
{
	if(ElectricalCharge(speciesInfo.formula) != 0)
		return false;
	
	for(unsigned i = 0; i < reactionInfo.products.size(); ++i)
		if(ElectricalCharge(reactionInfo.products[i]) == 0 || reactionInfo.stoichiometries[i] < 0.0)
			return false;
	
	return true;
}

void AqueousPhase::InitElectrolytes()
{
	// Get a reference to the database object constaining species information
	const Database& db = Database::GetInstance();
	
	vector<vector<Index>>  ionsIndexes;
	vector<vector<double>> ionsStoichiometries;
	
	// Loop over all the aqueous species to determine the electrolytes
	for(unsigned ispecies = 0; ispecies < species.size(); ++ispecies)
	{
		const SpeciesInfo speciesInfo = db.GetSpeciesInfo(species[ispecies]);
		
		if(speciesInfo.HasTag("basis"))
			continue;
		
		const ReactionInfo reactionInfo = db.GetReactionInfo(species[ispecies]);
		
		if(!IsElectrolyte(speciesInfo, reactionInfo))
			continue;
		
		ionsIndexes.push_back(vector<Index>());
		ionsStoichiometries.push_back(vector<double>());
		
		electrolytes.push_back(ispecies);
		
		for(unsigned i = 0; i < reactionInfo.products.size(); ++i)
		{
			const Index ionGlobalIndex = std::find(species.begin(), species.end(), reactionInfo.products[i]) - species.begin();
			
			ionsIndexes.back().push_back( std::find(ions.begin(), ions.end(), ionGlobalIndex) - ions.begin() );
			
			ionsStoichiometries.back().push_back(reactionInfo.stoichiometries[i]);
		}
	}
	
	// Initialise the Eigen matrix to the dimension (numElectrolytes, numIons) with zero values
	vei = MatrixXd::Zero(electrolytes.size(), ions.size());
	
	for(unsigned e = 0; e < electrolytes.size(); ++e)
		for(unsigned iter = 0; iter < ionsIndexes[e].size(); ++iter)
			vei(e, ionsIndexes[e][iter]) = ionsStoichiometries[e][iter];
}

void AqueousPhase::InitDefaultActivityModels()
{
	// Setting the activity model for the ionic species in the aqueous phase
	ionicActivityFunc = HkfIonicActivityModel(GetIonicSpecies());
	
	// Setting the activity model for the neutral species in the aqueous phase
	neutralActivityFuncs.resize(neutrals.size());
	
	for(unsigned ineutral = 0; ineutral < neutrals.size(); ++ineutral)
	{
		if("H2O" == species[neutrals[ineutral]])
			neutralActivityFuncs[ineutral] = HkfWaterActivityModel(GetIonicSpecies());
		else
		if("CO2(aq)" == species[neutrals[ineutral]])
			neutralActivityFuncs[ineutral] = DrummondCO2ActivityModel();
		else
			neutralActivityFuncs[ineutral] = SetschenowActivityModel();
	}
}

const vector<double> AqueousPhase::ActivityCoefficientsIonicSpecies(double T, double P, const vector<double>& n) const
{
	return ionicActivityFunc(*this, T, P, n);
}

const vector<double> AqueousPhase::ActivityCoefficientsNeutralSpecies(double T, double P, const vector<double>& n) const
{
	// The activity coefficients of the neutral species
	vector<double> gn(neutrals.size());
	
	// Loop over all the neutral species and compute their activity coefficients 
	for(unsigned ineutral = 0; ineutral < neutrals.size(); ++ineutral)
		gn[ineutral] = neutralActivityFuncs[ineutral](*this, T, P, n);
	
	return gn;
}

